package cn.net.autocode.files.pdf;

import cn.net.autocode.files.ocr.OcrFactory;
import net.sourceforge.tess4j.Tesseract;
import net.sourceforge.tess4j.TesseractException;
import org.apache.pdfbox.Loader;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.ImageType;
import org.apache.pdfbox.rendering.PDFRenderer;

import javax.imageio.IIOImage;
import javax.imageio.ImageIO;
import javax.imageio.ImageWriteParam;
import javax.imageio.ImageWriter;
import javax.imageio.stream.MemoryCacheImageOutputStream;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.*;
import java.util.Iterator;

public class PdfOcrTools{

    public static String ocrParse(File file){
        try {
            InputStream inputStream = new FileInputStream(file);
            return PdfOcrTools.ocrParse(inputStream,null,144);
        } catch (FileNotFoundException e) {
            return e.getMessage();
        }
    }

    public static String ocrParse(File file,int dpi){
        try {
            InputStream inputStream = new FileInputStream(file);
            return PdfOcrTools.ocrParse(inputStream,null,dpi);
        } catch (FileNotFoundException e) {
            return e.getMessage();
        }
    }

    public static String ocrParse(File file,String language){
        try {
            InputStream inputStream = new FileInputStream(file);
            return PdfOcrTools.ocrParse(inputStream,language,144);
        } catch (FileNotFoundException e) {
            return e.getMessage();
        }
    }

    public static String ocrParse(File file,String language,int dpi){
        try {
            InputStream inputStream = new FileInputStream(file);
            return PdfOcrTools.ocrParse(inputStream,language,dpi);
        } catch (FileNotFoundException e) {
            return e.getMessage();
        }
    }


    public static String ocrParse(InputStream inputStream){
        return PdfOcrTools.ocrParse(inputStream,null,144);
    }

    public static String ocrParse(InputStream inputStream,int dpi){
        return PdfOcrTools.ocrParse(inputStream,null,dpi);
    }

    public static String ocrParse(InputStream inputStream,String language){
        return PdfOcrTools.ocrParse(inputStream,language,144);
    }


    public static String ocrParse(InputStream inputStream,String language,int dpi){
        if(language==null || language.isBlank()){
            language = "chi_sim";
        }
        StringBuilder strBuilder = new StringBuilder();
        try {
            PDDocument pdfDocument = Loader.loadPDF(inputStream.readAllBytes());
            PDFRenderer pdfRenderer = new PDFRenderer(pdfDocument);
            Tesseract tesseract = OcrFactory.createTesseract(language);
            for(int i=0;i<pdfDocument.getNumberOfPages();i++){
                BufferedImage bufferedImage = pdfRenderer.renderImageWithDPI(i, dpi, ImageType.RGB);
                String ocrResult = tesseract.doOCR(bufferedImage);
                strBuilder.append(ocrResult);
            }
        } catch (IOException | TesseractException e) {
            return e.getMessage();
        }
        // 输出识别结果
        return strBuilder.toString();
    }

    public static byte[] pdfToImageByteArray(InputStream inputStream, int dpi){
        try {
            PDDocument pdfDocument = Loader.loadPDF(inputStream.readAllBytes());
            PDFRenderer pdfRenderer = new PDFRenderer(pdfDocument);

            int pageCount = pdfDocument.getNumberOfPages();

            // 计算总高度和最大宽度
            int totalHeight = 0;
            int maxWidth = 0;
            BufferedImage[] pageImages = new BufferedImage[pageCount];
            for (int page = 0; page < pageCount; page++) {
                pageImages[page] = pdfRenderer.renderImageWithDPI(page, dpi, ImageType.RGB);
                totalHeight += pageImages[page].getHeight();
                maxWidth = Math.max(maxWidth, pageImages[page].getWidth());
            }

            // 创建长图片    `
            BufferedImage longImage = new BufferedImage(maxWidth, totalHeight, BufferedImage.TYPE_INT_RGB);
            Graphics2D g2d = longImage.createGraphics();
            g2d.setColor(Color.WHITE);
            g2d.fillRect(0, 0, maxWidth, totalHeight);

            // 拼接图片
            int currentY = 0;
            for (BufferedImage pageImage : pageImages) {
                g2d.drawImage(pageImage, 0, currentY, null);
                currentY += pageImage.getHeight();
            }
            g2d.dispose();

            // 将长图片转换为 byte[]
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            Iterator<ImageWriter> writers = ImageIO.getImageWritersByFormatName("png");
            if (!writers.hasNext()) {
                throw new IllegalStateException("No PNG image writer found.");
            }
            ImageWriter writer = writers.next();
            ImageWriteParam param = writer.getDefaultWriteParam();
            writer.setOutput(new MemoryCacheImageOutputStream(baos));
            writer.write(null, new IIOImage(longImage, null, null), param);
            writer.dispose();
            return baos.toByteArray();
        } catch (IOException e) {
            throw new RuntimeException(e.getMessage());
        }
    }
}
